/*
 * SPDX-FileCopyrightText: 2025 UnionTech Software Technology Co., Ltd.
 *
 * SPDX-License-Identifier: GPL-2.0-or-later
 */

use anyhow::{Context, Result};
use std::fs::File;
use std::io;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;

use crate::cli::ZipArgs;
use crate::error::ZipError;
use crate::utils::common::caculate_ratio;
use crate::utils::log::LogConfig;
use crate::zip::{
    CentralDirectoryHeader, CompressionMethod, LocalFileHeader, ZipArchive, ZipWriter, VERSION_MADE,
};

struct EntryInfo {
    start_offset: u64,
    end_offset: u64,
    local_header: LocalFileHeader,
}

// 修复模式枚举
enum FixMode {
    Normal, // 普通修复模式 zip -F
    Full,   // 全盘修复模式 zip -FF
}

pub fn execute(args: &ZipArgs) -> Result<()> {
    let zipfile = args
        .zipfile
        .as_ref()
        .ok_or_else(|| ZipError::InvalidArguments("No zip file specified".to_string()))?;

    if !zipfile.exists() {
        return Err(ZipError::ArchiveNotFound(zipfile.clone()).into());
    }

    let output_path = args
        .other
        .out
        .as_ref()
        .ok_or_else(|| ZipError::InvalidArguments("No output file specified".to_string()))?;

    // 解析修复模式
    let fix_mode = if args.fix.fix_normal {
        FixMode::Normal
    } else {
        FixMode::Full
    };

    match fix_mode {
        FixMode::Normal => {
            LogConfig::println("Fix archive (-F) - assume mostly intact archive");

            log::info!("use normal mode to repair zip ...");
            repair_zip_normal(zipfile, output_path, args)?;
        }
        FixMode::Full => {
            LogConfig::println("Fix archive (-FF) - salvage what can");

            log::info!("use full mode to repair zip...");
            repair_zip_full(zipfile, output_path, args)
                .with_context(|| "Failed to repair zip in full mode")?;
        }
    }
    log::info!("Repair completed successfully.");
    Ok(())
}

fn scan_local_file_headers(file: &mut File) -> anyhow::Result<Vec<(u32, u32, String)>> {
    let mut positions: Vec<(u32, u32, String)> = Vec::new();
    let file_size = file.metadata()?.len();
    let mut pos = 0u64;

    while pos + 4 < file_size {
        file.seek(SeekFrom::Start(pos))?;
        let mut sig = [0u8; 4];
        file.read_exact(&mut sig)?;
        if sig == [0x50, 0x4B, 0x03, 0x04] {
            // 读取文件名长度
            let mut buf = [0u8; 26];
            if file.read_exact(&mut buf).is_err() {
                pos += 1;
                continue;
            }
            let compressed_size = u32::from_le_bytes([buf[14], buf[15], buf[16], buf[17]]);
            let filename_len = u16::from_le_bytes([buf[22], buf[23]]) as usize;
            let extra_field_len = u16::from_le_bytes([buf[24], buf[25]]) as usize;

            // 读取文件名和额外字段
            let mut filename = vec![0u8; filename_len];
            file.read_exact(&mut filename)?;

            let end =
                pos as u32 + 30 + filename_len as u32 + extra_field_len as u32 + compressed_size;
            positions.push((
                pos as u32,
                end,
                String::from_utf8_lossy(&filename).to_string(),
            ));
            // 跳到下一个可能的文件头
            pos = file.stream_position()?;
        } else {
            pos += 1;
        }
    }
    Ok(positions)
}

/// zip -F：仅修复中央目录，跳过损坏条目
fn repair_zip_normal(input_path: &Path, output_path: &Path, _args: &ZipArgs) -> Result<()> {
    let mut file = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(input_path)
        .with_context(|| format!("Failed to open existing zip file: {}", input_path.display()))?;
    let mut writer = ZipWriter::new(output_path.to_str().unwrap())?;

    let file_size = file.metadata()?.len();
    let mut pos;

    // 优先查找ZIP结束签名
    if ZipArchive::find_end_record_signature(file.try_clone()?, file_size).is_err() {
        LogConfig::println_warning("bad archive - missing end signature");
        LogConfig::println_warning("Can't use -F to fix (try --FF)");
        LogConfig::println("");
        return Err(ZipError::InvalidArchive(input_path.to_str().unwrap().to_string()).into());
    } else {
        LogConfig::println("Zip entry offsets do not need adjusting");
    }

    let mut cd_headers = Vec::new();
    pos = 0;
    // 扫描文件查找中央目录
    while pos + 4 < file_size {
        file.seek(SeekFrom::Start(pos))?;
        let mut sig = [0u8; 4];
        file.read_exact(&mut sig)?;
        if sig == [0x50, 0x4B, 0x01, 0x02] {
            // 中央目录
            LogConfig::println_verbose(&format!(
                "central directory header signature on disk {} at {}",
                0, pos
            ));
            // 读取固定长度的头部字段
            let mut header_bytes = [0u8; 42];
            file.read_exact(&mut header_bytes)?;

            let version_made = u16::from_le_bytes([header_bytes[0], header_bytes[1]]);
            let version_needed = u16::from_le_bytes([header_bytes[2], header_bytes[3]]);
            let flags = u16::from_le_bytes([header_bytes[4], header_bytes[5]]);
            let compression = u16::from_le_bytes([header_bytes[6], header_bytes[7]]);
            let mod_time = u16::from_le_bytes([header_bytes[8], header_bytes[9]]);
            let mod_date = u16::from_le_bytes([header_bytes[10], header_bytes[11]]);
            let crc32 = u32::from_le_bytes([
                header_bytes[12],
                header_bytes[13],
                header_bytes[14],
                header_bytes[15],
            ]);
            let compressed_size = u32::from_le_bytes([
                header_bytes[16],
                header_bytes[17],
                header_bytes[18],
                header_bytes[19],
            ]);
            let uncompressed_size = u32::from_le_bytes([
                header_bytes[20],
                header_bytes[21],
                header_bytes[22],
                header_bytes[23],
            ]);
            let filename_len = u16::from_le_bytes([header_bytes[24], header_bytes[25]]);
            let extra_field_len = u16::from_le_bytes([header_bytes[26], header_bytes[27]]);
            let file_comment_len = u16::from_le_bytes([header_bytes[28], header_bytes[29]]);
            let disk_num = u16::from_le_bytes([header_bytes[30], header_bytes[31]]);
            let internal_attr = u16::from_le_bytes([header_bytes[32], header_bytes[33]]);
            let external_attr = u32::from_le_bytes([
                header_bytes[34],
                header_bytes[35],
                header_bytes[36],
                header_bytes[37],
            ]);
            let local_header_offset = u32::from_le_bytes([
                header_bytes[38],
                header_bytes[39],
                header_bytes[40],
                header_bytes[41],
            ]);

            // 读取文件名、额外字段和注释
            let mut filename = vec![0u8; filename_len as usize];
            file.read_exact(&mut filename)?;

            let mut extra_field = vec![0u8; extra_field_len as usize];
            file.read_exact(&mut extra_field)?;

            let mut file_comment = vec![0u8; file_comment_len as usize];
            file.read_exact(&mut file_comment)?;

            cd_headers.push(CentralDirectoryHeader {
                version_made,
                version_needed,
                flags,
                compression: CompressionMethod::from(compression),
                mod_time,
                mod_date,
                crc32,
                compressed_size,
                uncompressed_size,
                filename,
                extra_field,
                file_comment,
                disk_num,
                internal_attr,
                external_attr,
                local_header_offset,
                zip64_extended_info: None,
            });
            pos = file.stream_position()?;
        } else {
            pos += 1;
            continue;
        }
    }

    if cd_headers.is_empty() {
        log::warn!("bad archive - no central directory found");
        return Err(ZipError::InvalidArchive(input_path.to_str().unwrap().to_string()).into());
    }

    // 扫描本地文件头位置
    let local_headers = scan_local_file_headers(&mut file)?;

    let mut total_bytes: u64 = 0;
    let mut compressed_bytes: u64 = 0;
    // 遍历中央目录，并校验每个条目是否存在
    for header in cd_headers.iter_mut() {
        let file_name = String::from_utf8_lossy(&header.filename);
        LogConfig::println(&format!(" copying: {}", file_name));
        let mut data_end = 0u32;
        // 更新中央目录记录的本地文件头偏移量
        if let Some((offset, end, _)) = local_headers
            .iter()
            .find(|(_, _, fname)| fname == &file_name.to_string())
        {
            header.local_header_offset = *offset;
            data_end = *end;
        }

        // 验证本地文件和中央目录记录的是否匹配
        if ZipArchive::validate_local_header(&mut file, header).is_err() {
            LogConfig::println_warning(&format!("bad - skipping: {}", file_name));
            continue;
        }

        // 写入修复后的中央目录记录
        // 分块读取和写入数据
        let start = header.local_header_offset;
        let local_data_start = writer.get_file_stream_position()?;
        header.local_header_offset = local_data_start as u32;

        copy_data(&mut file, &mut writer, start, data_end)?;

        writer.add_central_directory_header(header.clone());

        total_bytes += header.uncompressed_size as u64;
        compressed_bytes += header.compressed_size as u64;
    }
    writer.finish()?;

    LogConfig::println_verbose(&format!(
        "total bytes={}, compressed={}, ratio={:.0}% savings",
        total_bytes,
        compressed_bytes,
        caculate_ratio(total_bytes, compressed_bytes)
    ));

    Ok(())
}

/// zip -FF：全盘扫描本地文件头
fn repair_zip_full(input_path: &Path, output_path: &Path, args: &ZipArgs) -> Result<()> {
    let mut input = File::open(input_path)?;
    let file_size = input.metadata()?.len();
    let mut pos = 0u64;
    let mut entries = Vec::new();
    let mut central_entries = Vec::new();

    // 是否找到文件尾标识
    let eocdr_pos = match ZipArchive::find_end_record_signature(input.try_clone()?, file_size) {
        Ok(pos) => {
            LogConfig::println_verbose(
                " Found end record (EOCDR) - says expect single disk archive",
            );
            pos
        }
        Err(_) => {
            LogConfig::println_warning(
                "Missing end (EOCDR) signature - either this archive\n                     is not readable or the end is damaged",
            );
            // 新增：用户确认是否为单盘
            print!("Is this a single-disk archive? (y/n): ");
            io::stdout().flush().unwrap();
            let mut input_str = String::new();
            io::stdin().read_line(&mut input_str).unwrap();
            let input_str = input_str.trim();
            if input_str != "y" && input_str != "Y" {
                LogConfig::println_warning("User chose not to continue. Abort.");
                return Err(anyhow::anyhow!(
                    "User aborted repair (not single-disk archive)"
                ));
            }
            LogConfig::println("  Assuming single-disk archive");
            0
        }
    };

    LogConfig::println("Scanning for entries...");

    // 扫描文件查找有效的本地文件头
    while pos + 4 < file_size {
        input.seek(SeekFrom::Start(pos))?;
        let mut sig = [0u8; 4];
        input.read_exact(&mut sig)?;

        if sig == [0x50, 0x4B, 0x03, 0x04] {
            //文件头
            let entry_start = pos;

            // 读取本地文件头
            let mut header = [0u8; 26];
            input.read_exact(&mut header)?;

            let version_needed = u16::from_le_bytes([header[0], header[1]]);
            let flags = u16::from_le_bytes([header[2], header[3]]);
            let compression = u16::from_le_bytes([header[4], header[5]]);
            let mod_time = u16::from_le_bytes([header[6], header[7]]);
            let mod_date = u16::from_le_bytes([header[8], header[9]]);
            let crc32 = u32::from_le_bytes([header[10], header[11], header[12], header[13]]);
            let compressed_size =
                u32::from_le_bytes([header[14], header[15], header[16], header[17]]);
            let uncompressed_size =
                u32::from_le_bytes([header[18], header[19], header[20], header[21]]);
            let filename_len = u16::from_le_bytes([header[22], header[23]]);
            let extra_field_len = u16::from_le_bytes([header[24], header[25]]);

            // 读取文件名和额外字段
            let mut filename = vec![0u8; filename_len as usize];
            let mut extra_field = vec![0u8; extra_field_len as usize];
            input.read_exact(&mut filename)?;
            input.read_exact(&mut extra_field)?;

            let data_offset = input.stream_position()?;
            let entry_end = data_offset + compressed_size as u64;

            LogConfig::print_verbose(&format!(" Local ({:>2} {:>6}):", 1, pos));
            LogConfig::print(&format!(
                " copying: {}  ({} bytes)",
                String::from_utf8_lossy(&filename),
                compressed_size
            ));
            if args.basic_options.verbose {
                LogConfig::println_verbose(&format!("({} bytes)", compressed_size));
            } else {
                LogConfig::println("");
            }
            let local_file_entry = LocalFileHeader {
                version_needed,
                flags,
                compression: CompressionMethod::from(compression),
                mod_time,
                mod_date,
                crc32,
                compressed_size,
                uncompressed_size,
                filename,
                extra_field,
                zip64_extended_info: None,
            };

            entries.push(EntryInfo {
                local_header: local_file_entry,
                start_offset: entry_start,
                end_offset: entry_end,
            });

            pos = entry_end;
        } else if sig == [0x50, 0x4B, 0x01, 0x02] {
            //中央目录
            if central_entries.is_empty() {
                LogConfig::println("Central Directory found...");
            }
            LogConfig::print_verbose(&format!(" {:<5} ({:>2} {:>6}): ", "Cen", 1, pos,));
            let mut header = [0u8; 42];
            if input.read_exact(&mut header).is_err() {
                log::warn!("read next central directory header error");
                LogConfig::println("");
                LogConfig::println_warning("reading central directory: Invalid argument");
                LogConfig::println_warning("bad archive - error reading central directory");
                LogConfig::println_warning("skipping this entry...");
                pos += 1;
                continue;
            };

            // 解析中央目录头
            let version_made = u16::from_le_bytes([header[0], header[1]]);
            let version_needed = u16::from_le_bytes([header[2], header[3]]);
            let flags = u16::from_le_bytes([header[4], header[5]]);
            let compression = u16::from_le_bytes([header[6], header[7]]);
            let mod_time = u16::from_le_bytes([header[8], header[9]]);
            let mod_date = u16::from_le_bytes([header[10], header[11]]);
            let crc32 = u32::from_le_bytes([header[12], header[13], header[14], header[15]]);
            let compressed_size =
                u32::from_le_bytes([header[16], header[17], header[18], header[19]]);
            let uncompressed_size =
                u32::from_le_bytes([header[20], header[21], header[22], header[23]]);
            let filename_len = u16::from_le_bytes([header[24], header[25]]);
            let extra_field_len = u16::from_le_bytes([header[26], header[27]]);
            let file_comment_len = u16::from_le_bytes([header[28], header[29]]);
            let disk_num = u16::from_le_bytes([header[30], header[31]]);
            let internal_attr = u16::from_le_bytes([header[32], header[33]]);
            let external_attr =
                u32::from_le_bytes([header[34], header[35], header[36], header[37]]);
            let local_header_offset =
                u32::from_le_bytes([header[38], header[39], header[40], header[41]]);

            // 读取文件名、额外字段和注释
            let mut filename = vec![0u8; filename_len as usize];
            let mut extra_field = vec![0u8; extra_field_len as usize];
            let mut file_comment = vec![0u8; file_comment_len as usize];
            if input.read_exact(&mut filename).is_err() {
                LogConfig::println("");
                LogConfig::println_warning("reading central directory: Invalid argument");
                LogConfig::println_warning("error reading entry:  Invalid argument");
                LogConfig::println_warning("skipping this entry...");
                pos += 1;
                continue;
            }
            if input.read_exact(&mut extra_field).is_err() {
                LogConfig::println("");
                LogConfig::println_warning("error reading entry:  Invalid argument");
                LogConfig::println_warning("skipping this entry...");
                pos += 1;
                continue;
            }
            if input.read_exact(&mut file_comment).is_err() {
                LogConfig::println("");
                LogConfig::println_warning("error reading entry:  Invalid argument");
                LogConfig::println_warning("skipping this entry...");
                pos += 1;
                continue;
            }

            central_entries.push(CentralDirectoryHeader {
                version_made,
                version_needed,
                flags,
                compression: CompressionMethod::from(compression),
                mod_time,
                mod_date,
                crc32,
                compressed_size,
                uncompressed_size,
                filename: filename.clone(),
                extra_field,
                file_comment,
                disk_num,
                internal_attr,
                external_attr,
                local_header_offset,
                zip64_extended_info: None,
            });

            LogConfig::println_verbose(&format!(
                "updating: {}",
                String::from_utf8_lossy(&filename)
            ));

            pos = input.stream_position()?;
        } else {
            pos += 1;
        }
    }

    if eocdr_pos > 0 {
        LogConfig::println(&format!("EOCDR found ({:>2} {:>6})...", 1, eocdr_pos));
    }

    // 创建新的zip文件并写入找到的条目
    let mut writer = ZipWriter::new(output_path.to_str().unwrap())?;

    // 新增：如果没有修复任何条目才报错，否则返回Ok
    if entries.is_empty() {
        return Err(anyhow::anyhow!(
            "No entries could be repaired from the archive"
        ));
    }
    for entry in entries {
        // 写入本地文件头和数据
        let header_start = writer.get_file_stream_position()?;
        if copy_data(
            &mut input,
            &mut writer,
            entry.start_offset as u32,
            entry.end_offset as u32,
        )
        .is_err()
        {
            continue;
        }

        // 查找是否在中央目录中存在对应条目
        let cdh = if let Some(central_header) = central_entries
            .iter()
            .find(|ch| ch.filename == entry.local_header.filename)
        {
            // 使用中央目录中的条目，更新必要字段
            CentralDirectoryHeader {
                version_made: central_header.version_made,
                version_needed: central_header.version_needed,
                flags: central_header.flags,
                compression: central_header.compression,
                mod_time: central_header.mod_time,
                mod_date: central_header.mod_date,
                crc32: central_header.crc32,
                compressed_size: central_header.compressed_size,
                uncompressed_size: central_header.uncompressed_size,
                filename: central_header.filename.clone(),
                extra_field: central_header.extra_field.clone(),
                file_comment: central_header.file_comment.clone(),
                disk_num: central_header.disk_num,
                internal_attr: central_header.internal_attr,
                external_attr: central_header.external_attr,
                local_header_offset: header_start as u32, // 更新本地头偏移
                zip64_extended_info: None,
            }
        } else {
            // 如果没有找到对应的中央目录头，则使用本地文件头创建一个新的中央目录头
            CentralDirectoryHeader {
                version_made: VERSION_MADE,
                version_needed: entry.local_header.version_needed,
                flags: entry.local_header.flags,
                compression: entry.local_header.compression,
                mod_time: entry.local_header.mod_time,
                mod_date: entry.local_header.mod_date,
                crc32: entry.local_header.crc32,
                compressed_size: entry.local_header.compressed_size,
                uncompressed_size: entry.local_header.uncompressed_size,
                filename: entry.local_header.filename.clone(),
                extra_field: entry.local_header.extra_field.clone(),
                file_comment: Vec::new(),
                disk_num: 0,
                internal_attr: 0,
                external_attr: 0,
                local_header_offset: header_start as u32,
                zip64_extended_info: None,
            }
        };
        writer.add_central_directory_header(cdh);
    }
    writer.finish()?;
    Ok(())
}

/// 拷贝数据块从源文件到目标写入器
fn copy_data(file: &mut File, writer: &mut ZipWriter, start: u32, end: u32) -> Result<()> {
    let mut remaining = end.saturating_sub(start);
    let mut buffer = vec![0u8; 1024 * 1024]; // 1MB缓冲区
    file.seek(SeekFrom::Start(start as u64))?;

    while remaining > 0 {
        let chunk_size = std::cmp::min(remaining as usize, buffer.len());
        let chunk = &mut buffer[..chunk_size];

        file.read_exact(chunk)?;
        writer.inner_write_all(chunk)?;

        remaining -= chunk_size as u32;
    }
    Ok(())
}
